import { useState, useEffect, useRef } from 'react';
import axios from 'axios';
import {
  Layout,
  Typography,
  Card,
  Upload,
  Button,
  Input,
  Select,
  Alert,
  Divider,
  Space,
  Progress,
  notification,
  Statistic,
  Timeline,
  Tag,
  Modal
} from 'antd';
import { 
  UploadOutlined, 
  RobotOutlined, 
  SendOutlined,
  CheckCircleOutlined,
  SyncOutlined,
  ClockCircleOutlined,
  SettingOutlined
} from '@ant-design/icons';
import type { UploadFile, UploadProps } from 'antd';
import './App.css';
import ModelConfig from './components/ModelConfig';

const { Header, Content } = Layout;
const { Title, Paragraph, Text } = Typography;
const { TextArea } = Input;

// 后端服务地址
const API_BASE_URL = 'http://127.0.0.1:5001';

interface BuildProgress {
  active: boolean;
  stage: string; // idle | chunk | extract | embedding | done | error
  build_id: string | null;
  current_step: string;
  chunk_count: number;
  current_chunk_index: number;
  extract_progress: { total: number; done: number };
  embedding_progress: { total: number; done: number };
  entity_progress: { total: number; done: number; entities: number; relations: number };
  llm_calls: { total: number; done: number };
  updated_at: number;
  error_message: string;
}

function App() {
  // 文件上传状态
  const [fileList, setFileList] = useState<UploadFile[]>([]);
  const [isBuilding, setIsBuilding] = useState(false);
  const [buildProgress, setBuildProgress] = useState<BuildProgress | null>(null);
  
  // 模型配置对话框状态
  const [configModalVisible, setConfigModalVisible] = useState(false);
  
  // 构建进度详情
  const [timelineItems, setTimelineItems] = useState<any[]>([]);
  
  // 查询状态
  const [question, setQuestion] = useState('');
  const [queryMode, setQueryMode] = useState('local');
  const [answer, setAnswer] = useState('');
  const [isQuerying, setIsQuerying] = useState(false);
  
  // SSE连接引用
  const sseRef = useRef<EventSource | null>(null);
  
  // 清理SSE连接
  useEffect(() => {
    return () => {
      if (sseRef.current) {
        sseRef.current.close();
        sseRef.current = null;
      }
    };
  }, []);
  
  // 启动SSE连接监听构建进度
  const startBuildStream = (buildId: string) => {
    if (sseRef.current) {
      sseRef.current.close();
    }
    
    const eventSource = new EventSource(`${API_BASE_URL}/build/stream?build_id=${buildId}`);
    sseRef.current = eventSource;
    
    eventSource.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data);
        console.log('SSE消息:', data);
        
        // 更新进度状态
        if (data.progress) {
          setBuildProgress(data.progress);
        }
        
        // 更新时间线
        if (data.message) {
          const newItem = {
            children: data.message,
            dot: getStageIcon(data.stage),
            color: getStageColor(data.stage)
          };
          setTimelineItems(prev => [...prev, newItem]);
        }
        
        // 处理构建完成或错误
        if (data.stage === 'done') {
          setIsBuilding(false);
          notification.success({
            message: '构建完成',
            description: '知识图谱构建成功，现在可以开始问答了！'
          });
          eventSource.close();
          sseRef.current = null;
        } else if (data.stage === 'error') {
          setIsBuilding(false);
          notification.error({
            message: '构建失败',
            description: data.error || '构建过程中发生错误'
          });
          eventSource.close();
          sseRef.current = null;
        }
        
      } catch (error) {
        console.error('解析SSE消息失败:', error);
      }
    };
    
    eventSource.onerror = () => {
      console.error('SSE连接错误');
      setIsBuilding(false);
      eventSource.close();
      sseRef.current = null;
    };
  };
  
  // 获取阶段图标
  const getStageIcon = (stage: string) => {
    switch (stage) {
      case 'chunk':
        return <SyncOutlined spin />;
      case 'extract':
        return <SyncOutlined spin />;
      case 'embedding':
        return <SyncOutlined spin />;
      case 'done':
        return <CheckCircleOutlined />;
      case 'error':
        return <ClockCircleOutlined />;
      default:
        return <ClockCircleOutlined />;
    }
  };
  
  // 获取阶段颜色
  const getStageColor = (stage: string) => {
    switch (stage) {
      case 'chunk':
      case 'extract':
      case 'embedding':
        return 'blue';
      case 'done':
        return 'green';
      case 'error':
        return 'red';
      default:
        return 'gray';
    }
  };
  
  // 文件上传配置
  const uploadProps: UploadProps = {
    onRemove: () => {
      setFileList([]);
    },
    beforeUpload: (file) => {
      setFileList([file]);
      return false; // 阻止自动上传
    },
    fileList,
    maxCount: 1,
    accept: ".md"
  };
  
  // 处理构建请求
  const handleBuild = async () => {
    if (fileList.length === 0) {
      notification.warning({
        message: '请选择文件',
        description: '请先选择一个Markdown文件进行上传'
      });
      return;
    }
    
    const formData = new FormData();
    const file = fileList[0] as any;
    formData.append('file', file.originFileObj || file, file.name);
    
    setIsBuilding(true);
    setBuildProgress(null);
    setTimelineItems([]);
    
    try {
      const response = await axios.post(`${API_BASE_URL}/build`, formData, {
        headers: {
          'Content-Type': 'multipart/form-data',
        },
      });
      
      const buildId = response.data.build_id;
      if (buildId) {
        startBuildStream(buildId);
        notification.info({
          message: '开始构建',
          description: '文件上传成功，正在构建知识图谱...'
        });
      }
    } catch (error: any) {
      setIsBuilding(false);
      const errorMsg = error.response?.data?.error || error.message;
      notification.error({
        message: '构建失败',
        description: errorMsg
      });
    }
  };
  
  // 处理查询请求
  const handleQuery = async () => {
    if (!question.trim()) {
      notification.warning({
        message: '请输入问题',
        description: '请输入您要查询的问题'
      });
      return;
    }
    
    setIsQuerying(true);
    setAnswer('');
    
    try {
      const response = await axios.post(`${API_BASE_URL}/query`, {
        question: question.trim(),
        query_mode: queryMode,
        response_type: '要点式，尽量精炼'
      });
      
      setAnswer(response.data.answer);
    } catch (error: any) {
      const errorMsg = error.response?.data?.error || error.message;
      setAnswer(`查询失败: ${errorMsg}`);
      notification.error({
        message: '查询失败',
        description: errorMsg
      });
    } finally {
      setIsQuerying(false);
    }
  };
  
  // 获取进度百分比
  const getProgressPercent = (progress: { total: number; done: number }) => {
    if (progress.total === 0) return 0;
    return Math.round((progress.done / progress.total) * 100);
  };
  
  return (
    <Layout style={{ minHeight: '100vh', background: '#f5f5f5' }}>
      <Header style={{ 
        background: '#001529', 
        display: 'flex', 
        alignItems: 'center', 
        justifyContent: 'space-between',
        padding: '0 16px',
        boxShadow: '0 2px 8px rgba(0,0,0,0.15)',
        width: '100%'
      }}>
        <div style={{ display: 'flex', alignItems: 'center' }}>
          <RobotOutlined style={{ fontSize: '28px', color: 'white', marginRight: '12px' }} />
          <Title level={2} style={{ color: 'white', margin: 0 }}>
            GraphRAG 反应器知识平台
          </Title>
        </div>
        <Button 
          type="text" 
          icon={<SettingOutlined />} 
          style={{ color: 'white' }}
          onClick={() => setConfigModalVisible(true)}
        >
          模型设置
        </Button>
      </Header>
      
      <Content>
        <div className="main-flex-container">
          {/* 左侧：构建区域 */}
          <div className="flex-card-container">
            <Card 
              title="第一步：构建知识图谱" 
              style={{ 
                height: '100%',
                borderRadius: '8px',
                boxShadow: '0 2px 8px rgba(0,0,0,0.06)'
              }}
              headStyle={{ 
                background: '#fafafa',
                borderBottom: '1px solid #e8e8e8',
                borderRadius: '8px 8px 0 0'
              }}
            >
              <Space direction="vertical" style={{ width: '100%' }} size="middle">
                <Paragraph type="secondary">
                  上传一个 Markdown (.md) 格式的反应器设计文档，系统将自动解析并构建知识网络。
                </Paragraph>
                
                <Upload {...uploadProps}>
                  <Button icon={<UploadOutlined />}>选择 .md 文件</Button>
                </Upload>
                
                <Button
                  type="primary"
                  onClick={handleBuild}
                  loading={isBuilding}
                  disabled={fileList.length === 0}
                  size="large"
                  block
                >
                  {isBuilding ? '构建中...' : '开始构建'}
                </Button>
                
                {/* 构建进度显示 */}
                {buildProgress && buildProgress.active && (
                  <Card 
                    size="small" 
                    title="构建进度"
                    style={{ 
                      borderRadius: '6px',
                      border: '1px solid #d9d9d9'
                    }}
                  >
                    <Space direction="vertical" style={{ width: '100%' }}>
                      <div>
                        <Text strong>当前阶段: </Text>
                        <Tag color={getStageColor(buildProgress.stage)}>
                          {buildProgress.current_step}
                        </Tag>
                      </div>
                      
                      {buildProgress.chunk_count > 0 && (
                        <Statistic title="文档块数量" value={buildProgress.chunk_count} />
                      )}
                      
                      {buildProgress.extract_progress.total > 0 && (
                        <div>
                          <Space direction="vertical" style={{ width: '100%' }} size="small">
                            <div>
                              <Text>实体抽取进度 (chunk级):</Text>
                              <Progress 
                                percent={getProgressPercent(buildProgress.extract_progress)} 
                                status={buildProgress.stage === 'error' ? 'exception' : 'active'}
                                format={() => `${buildProgress.extract_progress.done}/${buildProgress.extract_progress.total}`}
                              />
                            </div>
                            
                            {buildProgress.stage === 'extract' && buildProgress.llm_calls.done > 0 && (
                              <div>
                                <Text>LLM调用统计:</Text>
                                <Tag color="blue">{buildProgress.llm_calls.done} 次调用</Tag>
                              </div>
                            )}
                            
                            {buildProgress.current_chunk_index > 0 && (
                              <Text type="secondary">
                                当前处理: chunk {buildProgress.current_chunk_index}/{buildProgress.chunk_count}
                              </Text>
                            )}
                          </Space>
                        </div>
                      )}
                      
                      {buildProgress.embedding_progress.total > 0 && (
                        <div>
                          <Text>向量化进度:</Text>
                          <Progress 
                            percent={getProgressPercent(buildProgress.embedding_progress)} 
                            status={buildProgress.stage === 'error' ? 'exception' : 'active'}
                            format={() => `${buildProgress.embedding_progress.done}/${buildProgress.embedding_progress.total}`}
                          />
                        </div>
                      )}
                      
                      {buildProgress.entity_progress.entities > 0 && (
                        <div style={{ display: 'flex', gap: '16px' }}>
                          <div style={{ flex: 1 }}>
                            <Statistic title="提取实体" value={buildProgress.entity_progress.entities} />
                          </div>
                          <div style={{ flex: 1 }}>
                            <Statistic title="提取关系" value={buildProgress.entity_progress.relations} />
                          </div>
                        </div>
                      )}
                    </Space>
                  </Card>
                )}
                
                {/* 构建时间线 */}
                {timelineItems.length > 0 && (
                  <Card 
                    size="small" 
                    title="处理日志"
                    style={{ 
                      borderRadius: '6px',
                      border: '1px solid #d9d9d9'
                    }}
                  >
                    <Timeline
                      items={timelineItems.slice(-5)} // 只显示最新5条
                      style={{ maxHeight: '200px', overflow: 'auto' }}
                    />
                  </Card>
                )}
                
                {buildProgress && buildProgress.error_message && (
                  <Alert
                    message="构建错误"
                    description={buildProgress.error_message}
                    type="error"
                    showIcon
                  />
                )}
              </Space>
            </Card>
          </div>
          
          {/* 右侧：查询区域 */}
          <div className="flex-card-container">
            <Card 
              title="第二步：开始问答" 
              style={{ 
                height: '100%',
                borderRadius: '8px',
                boxShadow: '0 2px 8px rgba(0,0,0,0.06)'
              }}
              headStyle={{ 
                background: '#fafafa',
                borderBottom: '1px solid #e8e8e8',
                borderRadius: '8px 8px 0 0'
              }}
            >
              <Space direction="vertical" style={{ width: '100%' }} size="middle">
                <Paragraph type="secondary">
                  构建完成后，您可以向知识图谱提问，系统会基于文档内容给出智能回答。
                </Paragraph>
                
                <TextArea
                  value={question}
                  onChange={(e) => setQuestion(e.target.value)}
                  placeholder="请输入您的问题，例如：这些文本涉及哪些反应器类型、关键设备与运行参数？"
                  rows={6}
                  size="large"
                />
                
                <div style={{ display: 'flex', gap: '16px', alignItems: 'center' }}>
                  <div style={{ flex: 1 }}>
                    <Space>
                      <Text>检索模式:</Text>
                      <Select
                        value={queryMode}
                        onChange={(value) => setQueryMode(value)}
                        style={{ width: 120 }}
                        options={[
                          { value: 'local', label: 'Local' },
                          { value: 'global', label: 'Global' },
                          { value: 'naive', label: 'Naive' }
                        ]}
                      />
                    </Space>
                  </div>
                  <div style={{ flex: 1 }}>
                    <Button
                      type="primary"
                      icon={<SendOutlined />}
                      onClick={handleQuery}
                      loading={isQuerying}
                      size="large"
                      block
                    >
                      提交问题
                    </Button>
                  </div>
                </div>
                
                {/* 回答显示区域 */}
                {(isQuerying || answer) && (
                  <>
                    <Divider>智能回答</Divider>
                    {isQuerying ? (
                      <div style={{ textAlign: 'center', padding: '20px' }}>
                        <SyncOutlined spin style={{ fontSize: '24px' }} />
                        <div style={{ marginTop: '10px' }}>正在思考中...</div>
                      </div>
                    ) : (
                      answer && (
                        <Card 
                          size="small"
                          style={{ 
                            borderRadius: '6px',
                            border: '1px solid #d9d9d9',
                            backgroundColor: '#fafafa'
                          }}
                        >
                          <pre style={{ 
                            whiteSpace: 'pre-wrap', 
                            fontFamily: 'inherit',
                            margin: 0,
                            fontSize: '14px',
                            lineHeight: '1.6'
                          }}>
                            {answer}
                          </pre>
                        </Card>
                      )
                    )}
                  </>
                )}
              </Space>
            </Card>
          </div>
        </div>
      </Content>
      
      {/* 模型配置对话框 */}
      <Modal
        title="模型配置"
        open={configModalVisible}
        onCancel={() => setConfigModalVisible(false)}
        footer={null}
        width={800}
        destroyOnClose
      >
        <ModelConfig onSaved={() => setConfigModalVisible(false)} />
      </Modal>
    </Layout>
  );
}

export default App;

